import asyncio
import json
import time

from datasets import load_dataset

from lagent.agents.stream import PLUGIN_CN, AsyncAgentForInternLM, AsyncMathCoder, get_plugin_prompt
from lagent.llms import INTERNLM2_META
from lagent.llms.lmdeploy_wrapper import AsyncLMDeployPipeline
from lagent.prompts.parsers import PluginParser

# 设置异步事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

# 初始化异步 LMDeploy 模型
model = AsyncLMDeployPipeline(
    path='/root/autodl-tmp/models/internlm2_5-1_8b-chat',  # 模型路径
    meta_template=INTERNLM2_META,  # 使用 InternLM2 的元模板
    model_name='internlm-chat',  # 模型名称
    tp=1,  # 张量并行度
    top_k=1,  # 只保留概率最高的候选词
    temperature=1.0,  # 采样温度
    stop_words=['<|im_end|>', '<|action_end|>'],  # 停止词
    max_new_tokens=1024,  # 最大生成长度
)

# ----------------------- 测试数学解题功能 -----------------------
print('-' * 80, 'interpreter', '-' * 80)

# 加载数学问题数据集
ds = load_dataset('lighteval/MATH', split='test')
# 从数据集中每隔1000个样本选择一个，总共选择1000个问题  
problems = [item['problem'] for item in ds.select(range(0, 5000, 1000))]

# 创建异步数学解题代理
coder = AsyncMathCoder(
    llm=model,
    interpreter=dict(
        type='lagent.actions.AsyncIPythonInterpreter',  # 使用异步 Python 解释器
        max_kernels=300  # 最大并发内核数
    ),
    max_turn=11  # 最大对话轮数
)

# 开始批量解题
tic = time.time()  # 记录开始时间
# 为每个问题创建一个异步任务
coros = [coder(query, session_id=i) for i, query in enumerate(problems)]
# 并发执行所有任务
res = loop.run_until_complete(asyncio.gather(*coros))
print('-' * 120)
print(f'time elapsed: {time.time() - tic}')  # 打印总耗时

# 保存解题步骤到文件
with open('./tmp_1.json', 'w') as f:
    json.dump([coder.get_steps(i) for i in range(len(res))],
              f,
              ensure_ascii=False,
              indent=4)

# ----------------------- 测试插件功能 -----------------------
print('-' * 80, 'plugin', '-' * 80)
# 配置 Arxiv 搜索插件
plugins = [dict(type='lagent.actions.AsyncArxivSearch')]
# 创建带插件的异步代理
agent = AsyncAgentForInternLM(
    llm=model,
    plugins=plugins,
    output_format=dict(
        type=PluginParser,
        template=PLUGIN_CN,  # 使用中文插件模板
        prompt=get_plugin_prompt(plugins)  # 生成插件提示信息
    ))

# 测试插件性能
tic = time.time()
# 创建50个相同的论文搜索请求
coros = [
    agent(query, session_id=i)
    for i, query in enumerate(['LLM智能体方向的最新论文有哪些？'] * 50)
]
# 并发执行所有请求
res = loop.run_until_complete(asyncio.gather(*coros))
print('-' * 120)
print(f'time elapsed: {time.time() - tic}')  # 打印总耗时
